// src/script/interpreter.rs

use std::collections::VecDeque;

use ripemd::{Digest, Ripemd160};
use secp256k1::{Message, PublicKey, Secp256k1, ecdsa};
use tracing::{debug, error, warn};

use crate::{
    CompactSignatureBytes, OpCode, Script, TransactionId,
    script::{Operand, ScriptError},
};

//
// 为了验证一笔交易，所有交易输入都必须验证。以下是输入验证过程：
//      找到 outpoint 引用的交易。
//      通过其在交易中的索引找到输出。
//      获取输出脚本。
//      将输出脚本附加到输入脚本。
//      在 Script 解释器上执行连接的脚本。
//      如果脚本成功，则交易有效。
// 这就是这个想法。当我们检查 Script 和交易数据结构时，它将成为一个具体的过程。
//

#[derive(Default)]
pub struct ScriptInterpreter {
    stack: VecDeque<Operand>, // Stack used to hold the data during script execution
}

impl ScriptInterpreter {
    pub fn new() -> Self {
        Self {
            stack: VecDeque::new(),
        }
    }

    // fn print_stack(&self) {
    //     tracing::trace!(">>>> Stack: {:?}\n", self.stack);
    // }

    pub fn execute_script(
        &mut self,
        script: &Script,
        utxo_txid: &TransactionId,
    ) -> Result<bool, ScriptError> {
        // tracing::trace!("Executing script: {:?}, utxo_txid: {:?}", script, utxo_txid);

        for opcode in script.as_opcodes() {
            // tracing::trace!("Executing opcode: {:?}", opcode);

            match opcode {
                OpCode::OpDup => self.do_dump()?,
                OpCode::OpHash160 => self.do_hash160()?,
                OpCode::OpEqualverify => self.do_equal_verify()?,
                OpCode::OpChecksig => self.do_check_signature(utxo_txid)?,
                OpCode::OpData(data) => {
                    self.stack.push_back(Operand::Data(data.clone()));
                }
                OpCode::OpSignature(sig) => {
                    self.stack
                        .push_back(Operand::Data(sig.as_bytes_ref().to_vec()));
                }
                OpCode::OpPubKey(pubkey) => {
                    self.stack
                        .push_back(Operand::Data(pubkey.as_bytes_ref().to_vec()));
                }
                OpCode::OpPubKeyHash(pub_key_hash) => {
                    self.stack
                        .push_back(Operand::Data(pub_key_hash.as_bytes_ref().to_vec()));
                }
                OpCode::OpBlockHeight(height) => {
                    debug!("Block height: {}", height);
                }
            }

            // tracing::trace!("Stack after executing. Length: {}", self.stack.len());
            // self.print_stack();
        }

        let len = self.stack.len();
        if len != 1 {
            error!("Invalid script! Stack size after execution is {}", len);
            Err(ScriptError::InvalidScript)
        } else {
            //  at here self.stack is not empty
            match self.stack.pop_back().unwrap() {
                Operand::Bool(result) => Ok(result),
                Operand::Data(data) => {
                    error!(
                        "Invalid script! After execution, the data in stack is {:?}",
                        data
                    );
                    Err(ScriptError::InvalidScript)
                }
            }
        }
    }

    fn do_dump(&mut self) -> Result<(), ScriptError> {
        // Duplicate the top element of the stack
        if let Some(top) = self.stack.back() {
            self.stack.push_back(top.clone());
            Ok(())
        } else {
            error!("[do_dump] Invalid script! Stack is empty");
            Err(ScriptError::StackUnderflow)
        }
    }

    fn do_hash160(&mut self) -> Result<(), ScriptError> {
        // Perform SHA256 followed by RIPEMD160 on the top of the stack
        if let Some(op_data) = self.stack.pop_back() {
            if let Operand::Data(data) = op_data {
                let ripemd160_hash = Ripemd160::digest(&data);
                self.stack.push_back(Operand::Data(ripemd160_hash.to_vec()));
                Ok(())
            } else {
                error!(
                    "[do_hash160] Invalid script! Stack contains invalid data. Stack size is {}",
                    self.stack.len()
                );
                Err(ScriptError::InvalidScript)
            }
        } else {
            error!("[do_hash160] Invalid script! Stack is empty");
            Err(ScriptError::StackUnderflow)
        }
    }

    fn do_equal_verify(&mut self) -> Result<(), ScriptError> {
        // Check if the top two stack items are equal
        if self.stack.len() < 2 {
            error!(
                "[do_equal_verify] Invalid script! The length({}) of Stack is too small!",
                self.stack.len()
            );
            return Err(ScriptError::StackUnderflow);
        }
        let a = self.stack.pop_back().unwrap();
        let b = self.stack.pop_back().unwrap();

        if a == b {
            Ok(())
        } else {
            warn!("[equal_verify] false");
            Err(ScriptError::EqualverifyFalse)
        }
    }

    fn do_check_signature(&mut self, utxo_txid: &TransactionId) -> Result<(), ScriptError> {
        if self.stack.len() < 2 {
            error!(
                "[do_check_signature] Invalid script! The length({}) of Stack is too small!",
                self.stack.len()
            );
            return Err(ScriptError::StackUnderflow);
        }
        let pubkey = self.stack.pop_back().unwrap();
        let signature = self.stack.pop_back().unwrap();

        if let Operand::Data(pubkey) = pubkey {
            if let Operand::Data(signature) = signature {
                let signature = CompactSignatureBytes::from_vec(signature).unwrap();
                if !self.verify_signature(utxo_txid, &signature, &pubkey) {
                    self.stack.push_back(Operand::Bool(false)); // the signature is invalid
                } else {
                    self.stack.push_back(Operand::Bool(true)); // the signature is valid
                }
                Ok(())
            } else {
                error!(
                    "[do_check_signature -> signature] Invalid script! Stack contains invalid data. Stack size is {}",
                    self.stack.len()
                );
                Err(ScriptError::InvalidScript)
            }
        } else {
            error!(
                "[do_check_signature -> pubkey] Invalid script! Stack contains invalid data. Stack size is {}",
                self.stack.len()
            );
            Err(ScriptError::InvalidScript)
        }
    }

    fn verify_signature(
        &self,
        msg: &TransactionId,
        signature: &CompactSignatureBytes,
        pubkey: &[u8],
    ) -> bool {
        let message = Message::from_digest(*msg.as_bytes_ref());
        let sig = ecdsa::Signature::from_compact(signature.as_bytes_ref())
            .expect("compact signatures are 64 bytes; DER signatures are 68-72 bytes");
        let public_key = PublicKey::from_slice(pubkey)
            .expect("public keys must be 33 or 65 bytes, serialized according to SEC 2");

        let secp = Secp256k1::verification_only();
        secp.verify_ecdsa(&message, &sig, &public_key).is_ok()
    }
}

#[cfg(test)]
mod tests {
    use std::vec;

    use secp256k1::rand::rngs::OsRng;
    use secp256k1::{Message, Secp256k1};

    use super::*;

    use crate::constants::BLOCK_CHIAN_VERSION;
    use crate::script::p2pkh::{create_tx_input_script, create_tx_output_script};
    use crate::transaction::{Transaction, TransactionIn, TransactionOut, TransactionOutPoint};
    use crate::{CompactSignatureBytes, OpCode, PublicKeyBytes, TransactionId, TransactionValue};

    #[test]
    fn test_script_with_invalid_signature() {
        // Invalid script that should fail
        let script = Script::new(vec![
            OpCode::OpData(vec![1, 2, 3, 4]), // Push some data
            OpCode::OpChecksig, // This should fail as we didn't push any valid data for signature verification
        ]);

        let mut interpreter = ScriptInterpreter::new();
        #[rustfmt::skip]
        let utxo_txid = [
            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
            17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
        ];
        let utxo_txid = TransactionId::new(utxo_txid);
        let result = interpreter.execute_script(&script, &utxo_txid);

        assert!(result.is_err());
        assert_eq!(result.err().unwrap(), ScriptError::StackUnderflow);
    }

    fn new_tx_out_point(tx_id: [u8; 32], index: u32) -> TransactionOutPoint {
        let tx_id = TransactionId::new(tx_id);
        TransactionOutPoint::new(tx_id, index)
    }

    fn new_tx_out(value: TransactionValue, pub_key: &PublicKeyBytes) -> TransactionOut {
        let script_pubkey = create_tx_output_script(pub_key);
        TransactionOut::new(value, script_pubkey)
    }

    fn new_tx_in(
        out_point: TransactionOutPoint,
        sig: CompactSignatureBytes,
        pub_key: PublicKeyBytes,
    ) -> TransactionIn {
        let script_signature = create_tx_input_script(sig, pub_key);
        TransactionIn::new_with_sequence(out_point, script_signature, 0xFFFFFFFF)
    }

    fn new_tx(inputs: Vec<TransactionIn>, outputs: Vec<TransactionOut>) -> Transaction {
        Transaction::new(BLOCK_CHIAN_VERSION, inputs, outputs, 0)
    }

    // Alice is payer, Bob is payee
    fn new_alice_to_bob_tx(value: TransactionValue, pub_key_payee: &PublicKey) -> Transaction {
        let inputs: Vec<TransactionIn> = Vec::new(); // no input

        let pub_key = PublicKeyBytes::new(pub_key_payee.serialize());
        let output = new_tx_out(value, &pub_key);
        let outputs: Vec<TransactionOut> = vec![output];

        new_tx(inputs, outputs)
    }

    // Bob is payer. Bob spend utxo that is from Alice
    fn new_bob_spend_utxo_tx(
        out_point: TransactionOutPoint,
        sig_payer: CompactSignatureBytes,
        pub_key_payer: &PublicKey,
    ) -> Transaction {
        let outputs: Vec<TransactionOut> = Vec::new(); // no output
        let pub_key = PublicKeyBytes::new(pub_key_payer.serialize());
        let input = new_tx_in(out_point, sig_payer, pub_key);
        let inputs: Vec<TransactionIn> = vec![input];

        new_tx(inputs, outputs)
    }

    #[test]
    fn test_spend_utxo() {
        bc_util::init_logger_for_test(None);

        let secp = Secp256k1::new();
        // let (alice_secret_key, alice_public_key) = secp.generate_keypair(&mut OsRng);
        let (bob_secret_key, bob_public_key) = secp.generate_keypair(&mut OsRng);

        let prev_tx = new_alice_to_bob_tx(0x12345678, &bob_public_key);
        let prev_tx_id = prev_tx.calc_tx_id();
        let prev_tx_id: [u8; 32] = prev_tx_id.as_bytes();
        let out_point_prev_tx = new_tx_out_point(prev_tx_id, 0);

        let message = Message::from_digest(prev_tx_id);
        let sig_payer = secp.sign_ecdsa(&message, &bob_secret_key);
        let sig_payer = sig_payer.serialize_compact();
        let sig_payer = CompactSignatureBytes::new(sig_payer);

        let current_tx = new_bob_spend_utxo_tx(out_point_prev_tx, sig_payer, &bob_public_key);

        let mut interpreter = ScriptInterpreter::new();
        let mut script = current_tx.get_input_script(0);
        assert_eq!(script.length(), 2);
        script.extend(prev_tx.get_output_script(0));
        assert_eq!(script.length(), 7);

        let prev_tx_id = TransactionId::new(prev_tx_id);
        let result = interpreter.execute_script(&script, &prev_tx_id);
        println!("result: {:?}", result);
        assert!(result.is_ok());
        assert!(result.unwrap());
    }
}

#[cfg(test)]
mod tests_learning {
    use secp256k1::hashes::{Hash, sha256};
    use secp256k1::rand::rngs::OsRng;
    use secp256k1::{Message, Secp256k1, ecdsa};

    #[test]
    fn test_secret_key_pub_key() {
        let (priv_key, public_key) = bc_util::generate_secp256k1_keypair();
        println!("Secret Key: {:?}", priv_key);
        println!("Public Key: {:?}", public_key);

        let sec_key = priv_key.secret_bytes();
        let pub_key = public_key.serialize();

        assert_eq!(32, sec_key.len());
        assert_eq!(33, pub_key.len());
    }

    #[test]
    fn test_secp256k1() {
        let secp = Secp256k1::new();
        let (priv_key, public_key) = secp.generate_keypair(&mut OsRng);
        println!("Secret Key: {:?}", priv_key);
        println!("Public Key: {:?}", public_key);

        let digest = sha256::Hash::hash("Hello World!".as_bytes());
        let message = Message::from_digest(digest.to_byte_array());

        println!("Message: {:?}", message);

        let signature = secp.sign_ecdsa(&message, &priv_key);
        println!("Signature: {:?}", signature);
        let sig = signature.serialize_compact();
        let sig = ecdsa::Signature::from_compact(&sig).unwrap();
        assert_eq!(sig, signature);

        let is_valid = secp.verify_ecdsa(&message, &signature, &public_key).is_ok();
        assert!(is_valid);

        let tampered_digest = sha256::Hash::hash("tampered".as_bytes());
        let tampered_message = Message::from_digest(tampered_digest.to_byte_array());
        assert!(
            secp.verify_ecdsa(&tampered_message, &signature, &public_key)
                .is_err()
        );

        let tampered_signature = secp.sign_ecdsa(&tampered_message, &priv_key);
        assert!(
            secp.verify_ecdsa(&tampered_message, &tampered_signature, &public_key)
                .is_ok()
        );
    }

    #[test]
    fn test_sha256() {
        let message = "Hello, World!";
        let digest = sha256::Hash::hash(message.as_bytes());
        let expected = "dffd6021bb2bd5b0af676290809ec3a53191dd81c7f70a4b28688a362182986f";

        assert_eq!(hex::encode(digest.to_byte_array()), expected);

        use sha2::Digest;
        let digest = sha2::Sha256::digest(message.as_bytes());
        assert_eq!(hex::encode(digest), expected);
    }
}
